import numpy as np
import pydicom
from getMask import get_lung_pipe_mask
import os
from skimage.morphology import  erosion, dilation


def dfs(grid, cur_i, cur_j):
    if grid[cur_i][cur_j] != 1:
        return 0
    grid[cur_i][cur_j] = 0
    ans = 1
    for di in range(-1, 2):
        for dj in range(-1, 2):
            next_i, next_j = cur_i + di, cur_j + dj
            ans += dfs(grid, next_i, next_j)
    return ans


def get_split_slices_num(split_num, num_pipe):
    two_index = 0
    if 2 in num_pipe:
        two_index = num_pipe.index(2)
    else:
        while num_pipe[two_index] == 1:
            two_index += 1

    if split_num == 1:
        return two_index
    else:
        four_index = two_index
        while num_pipe[four_index] == 2:
            four_index += 1
        return four_index


def maxNumberPipe(image):
    area = []
    for i, l in enumerate(image):
        for j, n in enumerate(l):
            num = dfs(image, i, j)
            if num != 0:
                area.append(num)

    if len(area) == 2:
        if area[0] * 10 < area[1] or area[1] * 10 < area[0]:
            return 1
        else:
            return 2
    else:
        return len(area)


# split_num: the number of pipe wanting, only accept 1, 2 or 4
# image: the original lung_pipe mask, should only include boolean
def get_split_mask(image, split_num=2):
    mask = get_lung_pipe_mask(image, draw_lung_pipe=False)
    lung_grow = dilation(mask, np.ones([8,8,8]))
    lung_grow = erosion(mask, np.ones([4,4,4]))

    if split_num == 4:
        return mask
    sum_length = len(image)
    slice_lung = np.zeros([sum_length, 512, 512])
    slice_lung[lung_grow] = 1
    num_pipe = []
    for i in range(sum_length):
        num_pipe.append(maxNumberPipe(slice_lung[i]))

    slice_num = get_split_slices_num(split_num, num_pipe)

    return np.append(slice_lung[0:slice_num], np.zeros([sum_length + 1 - slice_num, 512, 512]), axis=0)


if __name__ == '__main__':
    filepath = '/Users/kim/Downloads/2_97AAA705A8A44474BADECCF26BA35457/'
    count = 0
    image = []
    for file in os.listdir(filepath):
        child = os.path.join(filepath, file)
        if os.path.splitext(child)[-1] == ".dcm":
            # print(child)
            count += 1
            image.append(pydicom.read_file(child).pixel_array)

    print(len(image))
    np_img = np.array([image])
    get_split_mask(np_img)
